module MDRoutine

using Utils
using TopolParser
using Printf

function md_functions()
    while true
        println(h("\n                    ************ MD functions ************", fc="c"))
        println(" 0 Return")
        println(" 1 Traditional MD cascade")
        println(" 2 Generating template mdp files")
        println(" 3 Generating template top files")
        para = sel_fun(Int32)
        if para == 0
            return
        elseif para == 1
            println()
            println(" 0 Return")
            println(" 1 Run MD protocol for protein solution")
            println(" 2 Run MD protocol for self-molecule solution")
            input = sel_fun(Int32)
            if input == 0
                continue
            elseif input == 1
                request_hint("Input protein name, e.g., 1zni (protein file name should be the same except for extension): ")
                pro = get_file("", "pdb")
                request_hint("Input forcefield to use (default: amber99sb-ildn): ")
                ff = get_input("amber99sb-ildn")
                request_hint("Input topol file name to be generated (default: topol.top): ")
                top = get_input("topol.top")
                request_hint("Input em file name (default: em.mdp): ")
                em = get_input("em.mdp")
                if !isfile(em)
                    println("$em does not exist, will create it now.")
                    add_em_mdp()
                end
                request_hint("Input md file name (default: md.mdp): ")
                md = get_input("md.mdp")
                if !isfile(md)
                    println("$md does not exist, will create it now.")
                    add_md_mdp()
                end
                request_hint("Input number of kernels (default: 18): ")
                nkernel = get_input(18, Int32)
                if nkernel < 1 || nkernel > Sys.CPU_THREADS
                    nkernel = Int(Sys.CPU_THREADS / 2)
                    warning_hint("Warning: number of kernels not valid. Automatically set to $nkernel.\n")
                end
                run_protein_md(pro, ff, top, em, md, nkernel)
            end
        elseif para == 2
            generate_mdp()
        elseif para == 3
            generate_top()
        end
    end
end

function run_protein_md(pro, ff, top, em, md, nkernel)
    run(`gmx pdb2gmx -f $pro.pdb -o $pro.gro -p $top -i posre.itp -ignh -ff $ff -water tip3p -nobackup`)
    run(`gmx editconf -f $pro.gro -o $(pro)_box.gro -d 1 -nobackup`)
    run(`gmx solvate -cp $(pro)_box.gro -o $(pro)_sol.gro -p $top -nobackup`)
    run(`gmx grompp -f $em -c $(pro)_sol.gro -p $top -o em.tpr -maxwarn 99 -nobackup`)
    run(pipeline(`echo SOL`, `gmx genion -s em.tpr -p $top -o $(pro)_sol.gro -neutral -nobackup`))
    run(`gmx grompp -f $em -c $(pro)_sol.gro -p $top -o em.tpr -maxwarn 99 -nobackup`)
    run(`gmx mdrun -v -deffnm em -nt $nkernel -nobackup`)
    run(`gmx grompp -f $md -c em.gro -p $top -o md.tpr -maxwarn 99 -nobackup`)
    run(`gmx mdrun -v -deffnm md -nt $nkernel -nobackup`)
end

function generate_mdp()
    request_hint("\nSelect md mission type:\n")
    println(" 0 Energy minimization")
    println(" 1 Production simulation")
    input = sel_fun(Int32)
    if input == 0
        mdp_path = get_path("em.mdp", false)
        write(mdp_path, add_em_mdp())
    elseif input == 1
        mdp_path = get_path("md.mdp", false)
        write(mdp_path, add_md_mdp())
    end
    finish_path(mdp_path)
end

function generate_top()
    top_path = get_path("topol.top", false)
    write(top_path, gen_top_str())
    finish_path(top_path)
end

function add_em_mdp()::String
    request_hint("Select minimization algorithm:\n")
    println(" 0 Steepest Descent")
    println(" 1 Conjugate Gradient")
    algo_num = sel_fun(Int32)
    if algo_num == 0
        algo = "steep"
    elseif algo_num == 1
        algo = "cg"
    else
        error_hint("Error input.")
        return
    end
    request_hint("Input minimization steps (default: 2000): ")
    steps = get_input(2000, Int32)
    request_hint("Input minimization tolerance (default: 1000): ")
    tolerance = get_input(1000, Int32)
    s = "define = -DFLEXIBLE
integrator = $algo
nsteps = $steps
emtol  = $tolerance
emstep = 0.01
;
nstxout    = 100
nstlog     = 50
nstenergy  = 50
nstcgsteep = 100
;
pbc = xyz
cutoff-scheme            = Verlet
coulombtype              = PME
rcoulomb                 = 1.0
vdwtype                  = Cut-off
rvdw                     = 1.0
DispCorr                 = EnerPres
;
constraints              = none
;
"
    return s
end

function add_md_mdp()::String
    request_hint("Input dt (default: 0.002 ps): ")
    dt = get_input(0.002, Float64)
    request_hint("Input simulation time in unit ns (default: 40 ns): ")
    sim_time = get_input(40, Float64)
    request_hint("Input group name to eliminate center of mass (com) motation (default: protein): ")
    comm_grps = get_input("protein")
    request_hint("Eliminate angular also? [Y/n] ")
    comm_mode = get_input("y") != "n" ? "Angular" : ""
    request_hint("Input target temperature, unit: K (default: 298.15): ")
    ref_t = get_input("298.15")
    request_hint("Input target pressure, unit: bar (default: 1.0): ")
    ref_p = get_input("1.0")
    request_hint("Input group names to control temperature, splitted by space (default: protein Water_and_ions): ")
    tc_grps = get_input("protein Water_and_ions")
    num_grps = length(split(tc_grps))
    s = "define =
integrator = md
dt         = $dt   ; ps
nsteps     = $(convert(Int32, sim_time * 1000 / dt))   ; $(sim_time) ns
comm-grps  = $comm_grps
comm_mode  = $comm_mode
energygrps = 
;
nstxout = 0
nstvout = 0
nstfout = 0
nstlog  = 1000
nstenergy = 1000
nstxout-compressed = 1000
compressed-x-grps  = system
;
pbc = xyz
cutoff-scheme = Verlet
coulombtype   = PME
rcoulomb      = 1.0
vdwtype       = cut-off
rvdw          = 1.0
DispCorr      = EnerPres
;
Tcoupl  = V-rescale
tau_t   = $("0.2 " ^ num_grps)
tc_grps = $tc_grps
ref_t   = $("$ref_t " ^ num_grps)
;
Pcoupl     = parrinello-rahman
pcoupltype = isotropic
tau_p      = 2.0
ref_p      = $ref_p
compressibility = 4.5e-5
;
gen_vel  = no
gen_temp = 298.15
gen_seed = -1
;
freezegrps  = 
freezedim   = 
constraints = hbonds
;
"
    s = add_anneal_mdp(s, ref_t, num_grps)
    return s
end

function add_anneal_mdp(s::String, ref_t, num_grps)::String
    request_hint("Input simulated annealing type:\n")
    println(" 0 No simulated annealing settings")
    println(" 1 Single simulated annealing type")
    println(" 2 Periodic simulated annealing type")
    anneal_type = sel_fun(Int32)
    if anneal_type == 0
        return s
    elseif anneal_type == 1
        title_hint("System will be heated from 0 K at 0 ps to $ref_t K at 100 ps.\n")
        s = s * "annealing = $("single " ^ num_grps)
annealing_npoints = $("2 " ^ num_grps)
annealing_time = $("0 100 " ^ num_grps)
annealing_temp = $("0 $ref_t " ^ num_grps)
;
"
    elseif anneal_type == 2
        request_hint("Input the time points splitted with space (in unit ps): ")
        annealing_time = readline()
        request_hint("Input the temparature values splitted with space (in unit K): ")
        annealing_temp = readline()
        if length(split(annealing_time)) != length(split(annealing_temp))
            error_hint("Error: numbers of annealing time and temperature not match.\n")
            return s
        end
        annealing_npoints = length(split(annealing_time))
        s = s * "annealing = $("periodic " ^ num_grps)
annealing_npoints = $("$annealing_npoints " ^ num_grps)
annealing_time = $("$annealing_time " ^ num_grps)
annealing_temp = $("$annealing_temp " ^ num_grps)
;
"
        return s
    end
end

function gen_top_str()::String
    request_hint("Input forcefield name (default: amber99sb-ildn.ff): ")
    ff = get_input("amber99sb-ildn.ff")
    s = "; Example topology created by supernova (https://gitee.com/supernova_bingbing/SuperMDA)
;
; Typical top section sequence:
;
; 1. head       ; must be at the beginning of the topology file
; [ defaults ]
;
; 2. forcefield parameters                      ; must be located before '3. topology definations'
; [ atomtypes ]                                                                 ; included in ffnonbonded.itp
; [ bondtypes ] / [ constrainttypes ] / [ angletypes ] / [ dihedraltypes ]      ; included in ffbonded.itp
;
; 3. topology defination                        ; included in your own .itp file
; [ moleculetype ] / [ atoms ] / [ bonds ] / [ pairs ] / [ angles ] / [ dihedrals ] / [ exclusions ]
; [ position_restraints ]                       ; and other restraints 
; often one molecule per .itp file, e.g., protein, tip3p water, etc. (Repeatable)
;
; 4. system defination
; [ system ]
; [ molecules ]
;


; Include forcefield parameters
; If have own atom types, add them into ffnonbonded.itp, or include the own-type-containing .itp file firstly
#include \"$ff/forcefield.itp\"

"
    s, self_top_names = add_top_self(s)
    request_hint("\nInput water type (default: tip3p.itp): ")
    water_type = get_input("tip3p.itp")
    s = s * "\n; Include water topology
#include \"$ff/$water_type\"

#ifdef POSRES_WATER
; Position restraint for each water oxygen
[ position_restraints ]
;  i funct       fcx        fcy        fcz
1    1       1000       1000       1000
#endif

; Include topology for ions
#include \"$ff/ions.itp\"

"
    s = add_top_mol(s, self_top_names)
    return s
end

function add_top_self(s::String)
    request_hint("Input own topology file name (e.g., MOL.itp), ")
    println("one line per molecule.\nLeave blank if no self-topology or to stop:")
    self_top_files = Vector{AbstractString}()
    while true
        mol = get_input("")
        if isempty(mol)
            break
        end
        if !isfile(mol)
            error_hint("Not valid topol file.\n")
            continue
        end
        push!(self_top_files, mol)
    end
    
    s *= "; Include your molecule topology (containing POSRES defination if position-restraint simulation is to be done)\n"
    self_top_names = Vector{AbstractString}()
    for mol in self_top_files
        s *= "#include \"$mol\"\n"
        push!(self_top_names, gen_topol_from_file(mol).moleculetype)
    end
    return s, self_top_names
end

function add_top_mol(s::String, self_top_names::Vector{AbstractString})::String
    request_hint("Input system name (could be blank): ")
    s = s * "[ system ]
; Name
$(get_input("System generated by SuperMDA"))

[ molecules ]
; Compound        #mols
"
    while !isempty(self_top_names)
        mol = popfirst!(self_top_names)
        request_hint("Input number of $mol molecule (default: 1): ")
        nmol = get_input(1, Int32)
        s *= @sprintf("%s        %d\n", mol, nmol)
    end
    return s
end

end